import torch
import torch.nn as nn
# import torch.utils.data as data
from PIL import Image
import os
import numpy as np
from torchvision.transforms import transforms

from .dataset import ImageDataset

class ImageDataLoader(object):
    def __init__(self, train_image_path, test_img_path, img_path, train_mask_path, test_mask_path, train_batch_size, test_batch_size, n_threads):
        self.train_image_path = train_image_path
        self.test_image_path = test_img_path
        self.train_batch_size = train_batch_size
        self.test_batch_size = test_batch_size

        self.n_threads = n_threads
        self.train_mask_path = train_mask_path
        self.img_path = img_path
        self.test_mask_path = test_mask_path

        print("|===>Creating Image Data Loader")
        self.train_loader, self.test_loader = self.get_dataset()

    def get_dataset(self):
        normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                         std=[0.229, 0.224, 0.225])
        train_transforms = transforms.Compose([
            transforms.Resize(size=(512,512)),
            transforms.ToTensor(),
            normalize
        ])
        train_dataset = ImageDataset(self.train_image_path, self.train_mask_path, self.img_path, train_transforms, None)
        train_loader = torch.utils.data.DataLoader(
            train_dataset,
            batch_size = self.train_batch_size,
            shuffle= True,
            num_workers = self.n_threads,
            pin_memory = True
        )

        test_transforms = transforms.Compose([
            transforms.Resize(size=(512,512)),
            transforms.CenterCrop(size=(512,521)),
            transforms.ToTensor(),
            normalize
        ])
        test_dataset = ImageDataset(self.test_image_path, self.test_mask_path, self.img_path, test_transforms, None)

        test_loader = torch.utils.data.DataLoader(
            test_dataset,
            batch_size = self.test_batch_size,
            shuffle= True,
            num_workers = self.n_threads,
            pin_memory = True
        )
        return train_loader, test_loader

    def get_dataloader(self):
        return self.train_loader, self.test_loader


